"""
Nuclear Charge Radius Scaling Analysis
Author: Raheb Ali Mohammed Saleh Aoudh
Email: o.963852963852@gmail.com
Date: November 2024
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from typing import Dict, List, Tuple, Optional

class NuclearChargeRadiusAnalyzer:
    """
    A class to analyze lepton-mass-dependent effects in nuclear charge radius measurements.
    This implements the empirical scaling law derived from the proton radius puzzle.
    """
    
    def __init__(self):
        # Fundamental constants (CODATA 2014)
        self.m_e = 0.511  # MeV, electron mass
        self.m_mu = 105.658  # MeV, muon mass
        self.m_tau = 1777.0  # MeV, tau mass
        
        # Experimental data for nuclear charge radii
        self.nuclear_data = self._load_experimental_data()
        
        # Calculated parameters
        self.k_proton = None
        self.R0_proton = None
        self._calculate_proton_parameters()
    
    def _load_experimental_data(self) -> Dict:
        """
        Load experimental nuclear charge radius data from established measurements.
        
        Returns:
            Dictionary containing nuclear charge radius data
        """
        return {
            'proton': {
                'A': 1, 'Z': 1,
                'R_electron': 0.8751, 'R_electron_err': 0.0061,
                'R_muonic': 0.8409, 'R_muonic_err': 0.0004,
                'source': 'Pohl et al. 2010 + CODATA 2014'
            },
            'deuteron': {
                'A': 2, 'Z': 1,
                'R_electron': 2.125, 'R_electron_err': 0.003,
                'R_muonic': None, 'R_muonic_err': None,
                'source': 'Review of Modern Physics'
            },
            'helium-3': {
                'A': 3, 'Z': 2,
                'R_electron': 1.966, 'R_electron_err': 0.015,
                'R_muonic': None, 'R_muonic_err': None,
                'source': 'Nature Physics 2018'
            },
            'helium-4': {
                'A': 4, 'Z': 2,
                'R_electron': 1.681, 'R_electron_err': 0.004,
                'R_muonic': None, 'R_muonic_err': None,
                'source': 'Physical Review Letters 2020'
            }
        }
    
    def _calculate_proton_parameters(self):
        """
        Calculate the scaling parameters k and R0 from proton data.
        """
        R_e = self.nuclear_data['proton']['R_electron']
        R_mu = self.nuclear_data['proton']['R_muonic']
        
        # Calculate scaling parameter k
        mass_factor = (1/self.m_e - 1/self.m_mu)
        self.k_proton = (R_e - R_mu) / mass_factor
        
        # Calculate baseline radius R0
        self.R0_proton = R_e - self.k_proton / self.m_e
        
        # Calculate uncertainties
        k_err = self.nuclear_data['proton']['R_electron_err'] / mass_factor
        R0_err = np.sqrt(self.nuclear_data['proton']['R_electron_err']**2 + 
                        (k_err/self.m_e)**2)
        
        self.k_proton_err = k_err
        self.R0_proton_err = R0_err
    
    def scaling_law(self, m_lepton: float, R0: float, k: float) -> float:
        """
        The empirical scaling law: R(m) = R0 + k/m
        
        Args:
            m_lepton: Lepton mass in MeV
            R0: Baseline radius in fm
            k: Scaling parameter in fm·MeV
            
        Returns:
            Predicted charge radius in fm
        """
        return R0 + k / m_lepton
    
    def predict_muonic_radius(self, nucleus: str) -> Tuple[float, float]:
        """
        Predict muonic charge radius for a given nucleus.
        
        Args:
            nucleus: Name of the nucleus ('deuteron', 'helium-3', etc.)
            
        Returns:
            Tuple of (predicted radius, effective k parameter)
        """
        if nucleus not in self.nuclear_data:
            raise ValueError(f"Unknown nucleus: {nucleus}")
        
        data = self.nuclear_data[nucleus]
        A = data['A']
        
        # Calculate effective k for this nucleus
        k_effective = self.k_proton * A**(1/3)
        
        # Predict muonic radius
        R_mu_pred = data['R_electron'] - k_effective * (1/self.m_e - 1/self.m_mu)
        
        return R_mu_pred, k_effective
    
    def predict_tauonic_radius(self, nucleus: str) -> float:
        """
        Predict tauonic charge radius for a given nucleus.
        
        Args:
            nucleus: Name of the nucleus
            
        Returns:
            Predicted tauonic charge radius in fm
        """
        R_mu_pred, k_effective = self.predict_muonic_radius(nucleus)
        
        # For tauonic atoms, the difference from electronic is much smaller
        R_tau_pred = self.nuclear_data[nucleus]['R_electron'] - k_effective * (1/self.m_e - 1/self.m_tau)
        
        return R_tau_pred
    
    def calculate_statistical_significance(self, delta_R: float, 
                                         uncertainty: float = 0.005) -> float:
        """
        Calculate statistical significance of predicted difference.
        
        Args:
            delta_R: Predicted difference in fm
            uncertainty: Typical experimental uncertainty in fm
            
        Returns:
            Statistical significance in standard deviations
        """
        return abs(delta_R) / uncertainty
    
    def run_complete_analysis(self) -> pd.DataFrame:
        """
        Run complete analysis for all nuclei.
        
        Returns:
            DataFrame with complete analysis results
        """
        results = []
        
        print("Nuclear Charge Radius Scaling Analysis")
        print("=" * 50)
        print(f"Proton parameters: k = {self.k_proton:.6f} ± {self.k_proton_err:.6f} fm·MeV")
        print(f"Proton parameters: R0 = {self.R0_proton:.6f} ± {self.R0_proton_err:.6f} fm")
        print()
        
        for nucleus, data in self.nuclear_data.items():
            if nucleus == 'proton':
                # Use actual values for proton
                R_mu_pred = data['R_muonic']
                k_effective = self.k_proton
                delta_R = data['R_electron'] - data['R_muonic']
                significance = self.calculate_statistical_significance(delta_R)
                
                # Calculate R_tau for proton
                R_tau_pred = self.scaling_law(self.m_tau, self.R0_proton, self.k_proton)
                
            else:
                # Predict values for other nuclei
                R_mu_pred, k_effective = self.predict_muonic_radius(nucleus)
                delta_R = data['R_electron'] - R_mu_pred
                significance = self.calculate_statistical_significance(delta_R)
                R_tau_pred = self.predict_tauonic_radius(nucleus)
            
            result = {
                'nucleus': nucleus,
                'A': data['A'],
                'Z': data['Z'],
                'R_electron_fm': data['R_electron'],
                'R_muonic_predicted_fm': R_mu_pred,
                'R_tauonic_predicted_fm': R_tau_pred,
                'k_effective_fm_MeV': k_effective,
                'delta_R_fm': delta_R,
                'significance_sigma': significance,
                'source': data['source']
            }
            
            results.append(result)
            
            # Print results
            print(f"{nucleus.capitalize():<12} | A={data['A']} | "
                  f"R_e={data['R_electron']:.4f} fm | "
                  f"R_μ_pred={R_mu_pred:.4f} fm | "
                  f"ΔR={delta_R:.4f} fm | "
                  f"Significance={significance:.1f}σ")
        
        return pd.DataFrame(results)
    
    def plot_results(self, df: pd.DataFrame):
        """
        Create visualization of the analysis results.
        
        Args:
            df: DataFrame with analysis results
        """
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))
        
        # Plot 1: Electronic vs Muonic radii
        nuclei = df['nucleus']
        R_e = df['R_electron_fm']
        R_mu_pred = df['R_muonic_predicted_fm']
        
        x_pos = np.arange(len(nuclei))
        width = 0.35
        
        ax1.bar(x_pos - width/2, R_e, width, label='Electronic', alpha=0.7, color='blue')
        ax1.bar(x_pos + width/2, R_mu_pred, width, label='Muonic (Predicted)', alpha=0.7, color='red')
        ax1.set_xlabel('Nucleus')
        ax1.set_ylabel('Charge Radius (fm)')
        ax1.set_title('Electronic vs Predicted Muonic Charge Radii')
        ax1.set_xticks(x_pos)
        ax1.set_xticklabels(nuclei)
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        
        # Plot 2: Scaling parameter k vs Mass number A
        A_values = df['A']
        k_values = df['k_effective_fm_MeV']
        
        ax2.scatter(A_values, k_values, s=100, alpha=0.7, color='green')
        
        # Fit A^(1/3) curve
        A_fit = np.linspace(1, 4, 100)
        k_fit = self.k_proton * A_fit**(1/3)
        ax2.plot(A_fit, k_fit, 'r--', alpha=0.7, label='A^(1/3) fit')
        
        ax2.set_xlabel('Mass Number A')
        ax2.set_ylabel('k (fm·MeV)')
        ax2.set_title('Scaling Parameter k vs Mass Number')
        ax2.legend()
        ax2.grid(True, alpha=0.3)
        
        # Plot 3: Predicted differences
        delta_R = df['delta_R_fm']
        
        ax3.bar(nuclei, delta_R, alpha=0.7, color='orange')
        ax3.set_xlabel('Nucleus')
        ax3.set_ylabel('ΔR = R_e - R_μ (fm)')
        ax3.set_title('Predicted Differences Between Electronic and Muonic Radii')
        ax3.grid(True, alpha=0.3)
        
        # Plot 4: Statistical significance
        significance = df['significance_sigma']
        
        ax4.bar(nuclei, significance, alpha=0.7, color='purple')
        ax4.axhline(y=5, color='red', linestyle='--', label='5σ discovery threshold')
        ax4.set_xlabel('Nucleus')
        ax4.set_ylabel('Statistical Significance (σ)')
        ax4.set_title('Statistical Significance of Predicted Differences')
        ax4.legend()
        ax4.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.show()
    
    def uncertainty_analysis(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Perform uncertainty analysis on predictions.
        
        Args:
            df: DataFrame with analysis results
            
        Returns:
            DataFrame with uncertainty analysis
        """
        uncertainty_results = []
        
        for _, row in df.iterrows():
            nucleus = row['nucleus']
            data = self.nuclear_data[nucleus]
            
            # Uncertainty from electron radius measurement
            R_e_err = data['R_electron_err']
            
            # Uncertainty from k parameter
            k_err_prop = self.k_proton_err * data['A']**(1/3)
            
            # Total uncertainty in muonic prediction
            R_mu_err = np.sqrt(R_e_err**2 + (k_err_prop * (1/self.m_e - 1/self.m_mu))**2)
            
            # Relative uncertainty
            rel_uncertainty = (R_mu_err / row['R_muonic_predicted_fm']) * 100
            
            uncertainty_result = {
                'nucleus': nucleus,
                'R_electron_uncertainty_fm': R_e_err,
                'k_uncertainty_fm_MeV': k_err_prop,
                'R_muonic_uncertainty_fm': R_mu_err,
                'relative_uncertainty_percent': rel_uncertainty,
                'significance_with_uncertainty': row['delta_R_fm'] / R_mu_err
            }
            
            uncertainty_results.append(uncertainty_result)
        
        return pd.DataFrame(uncertainty_results)
    
    def export_results(self, df: pd.DataFrame, uncertainty_df: pd.DataFrame, 
                      filename: str = 'nuclear_radius_analysis.csv'):
        """
        Export analysis results to CSV file.
        
        Args:
            df: Main results DataFrame
            uncertainty_df: Uncertainty analysis DataFrame
            filename: Output filename
        """
        # Merge results
        full_results = pd.merge(df, uncertainty_df, on='nucleus')
        
        # Export to CSV
        full_results.to_csv(filename, index=False)
        print(f"\nResults exported to: {filename}")
        
        # Print summary
        print("\nAnalysis Summary:")
        print("=" * 40)
        for _, row in full_results.iterrows():
            print(f"{row['nucleus'].capitalize():<12}: ΔR = {row['delta_R_fm']:.4f} ± "
                  f"{row['R_muonic_uncertainty_fm']:.4f} fm "
                  f"({row['significance_with_uncertainty']:.1f}σ)")

def main():
    """
    Main function to run the complete analysis.
    """
    # Initialize analyzer
    analyzer = NuclearChargeRadiusAnalyzer()
    
    # Run complete analysis
    print("Starting Nuclear Charge Radius Scaling Analysis...")
    results_df = analyzer.run_complete_analysis()
    
    # Perform uncertainty analysis
    uncertainty_df = analyzer.uncertainty_analysis(results_df)
    
    # Create visualizations
    analyzer.plot_results(results_df)
    
    # Export results
    analyzer.export_results(results_df, uncertainty_df)
    
    # Print theoretical predictions
    print("\nTheoretical Predictions for Future Experiments:")
    print("=" * 50)
    for nucleus in ['deuteron', 'helium-3', 'helium-4']:
        R_mu_pred, k_eff = analyzer.predict_muonic_radius(nucleus)
        R_tau_pred = analyzer.predict_tauonic_radius(nucleus)
        print(f"{nucleus.capitalize():<12}: R_μ = {R_mu_pred:.4f} fm, "
              f"R_τ = {R_tau_pred:.4f} fm, k = {k_eff:.4f} fm·MeV")
    
    return analyzer, results_df, uncertainty_df

if __name__ == "__main__":
    analyzer, results, uncertainties = main()